#include "gemv.h"
#include <stdio.h>

#define WARP_SIZE 32
#define WARPS_PER_BLOCK 32
#define THREADS_PER_BLOCK 256  // WARP_SIZE * WARPS_PER_BLOCK


template <unsigned int WarpSize>
__device__ __forceinline__ float warpReduceSum(float sum) {
    if (WarpSize >= 32)sum += __shfl_down_sync(0xffffffff, sum, 16); // 0-16, 1-17, 2-18, etc.
    if (WarpSize >= 16)sum += __shfl_down_sync(0xffffffff, sum, 8);// 0-8, 1-9, 2-10, etc.
    if (WarpSize >= 8)sum += __shfl_down_sync(0xffffffff, sum, 4);// 0-4, 1-5, 2-6, etc.
    if (WarpSize >= 4)sum += __shfl_down_sync(0xffffffff, sum, 2);// 0-2, 1-3, 4-6, 5-7, etc.
    if (WarpSize >= 2)sum += __shfl_down_sync(0xffffffff, sum, 1);// 0-1, 2-3, 4-5, etc.
    return sum;
}

__global__ void fgemvKernel(float *A, float *B, float *C, size_t d0, size_t d1, size_t d4){


    size_t index = threadIdx.x + blockDim.x * blockIdx.x;
    if(index >= d0 * d1 * 8) return;

    int stride = 200;
    int b_offset = ((index / 8 )/stride) * 8 + (index % 8);

    float tmp = 0.0;
    for(int i =0; i < d4; i++){
        tmp += (A[index * d4  + i]) * (B[b_offset * d4 + i]);
    }

    C[index] = tmp;

}

//600*200*8*1*64 x 600*1*8*64*1 -> 600*200*8*1*1
__global__ void fgemvKernelV2(float *A, float *B, float *C, size_t d0, size_t d1, size_t d4){


    int bx = blockIdx.x;

    int tx = threadIdx.x;
    int ty = threadIdx.y;

    const int warp_size=32;
    int laneId= tx % warp_size;
    int current_row = blockDim.y * bx + ty;

    int stride = 200;

    int b_offset = ((current_row / 8 )/stride) * 8 + (current_row % 8);

    if(current_row < d0 * d1 * 8){
        float res=0;
        int kIteration = d4/warp_size;
        if(kIteration==0) kIteration=1;
        #pragma unroll
        for(int i=0; i< kIteration; i++){
            int current_col = i*warp_size + laneId;
            res += A[current_row*d4 + current_col] * B[b_offset * d4 + current_col];
        }
        res = warpReduceSum<warp_size>(res);
        if(laneId==0) C[current_row]=res;
    }

}

//600*200*8*1*64 x 600*1*8*64*1 -> 600*200*8*1*1
//d0 = 600
//d1 = 200
//d4 = 64
void fgemv(float *A, float *B, float *C, size_t d0, size_t d1, size_t d4, cudaStream_t stream) {

    dim3 block(THREADS_PER_BLOCK);
    dim3 grid(div_ceil(d0 * d1 * 8, THREADS_PER_BLOCK));

    //fgemvKernel<<<grid, block, 0, stream>>>(A, B, C, d0, d1, d4);

    fgemvKernelV2<<<d0 * d1 * 8/4, dim3{32,4},0, stream>>>(A, B, C, d0, d1, d4);
}